package org.zjx.server;

import lombok.Getter;
import lombok.Setter;
import org.zjx.client.SocketClient;
import org.zjx.core.SocketConfig;
import org.zjx.server.scanner.ServiceScanner;
import org.zjx.util.IpUtil;

import java.io.*;
import java.net.Socket;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.*;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@Getter
public class ClusterServer extends BaseSocketServer {
    private static final Logger logger = LoggerFactory.getLogger(ClusterServer.class);

    public enum ServerType {MASTER, WORKER}

    @Getter
    private final String serverId = UUID.randomUUID().toString();
    private final ServerType serverType;
    private RouteManager routeManager;
    private MasterConnector masterConnector;
    private final Map<String, SocketClient> workerConnections = new ConcurrentHashMap<>();
    private final ScheduledExecutorService backgroundExecutor = Executors.newScheduledThreadPool(2);
    private final ServiceRegistry serviceRegistry;

    @Getter
    @Setter
    private boolean running = false;

    public ClusterServer(SocketConfig config, ServerType serverType, String masterHost, int masterPort) {
        super(config);
        this.serverType = serverType;

        // 自动扫描并注册服务
        this.serviceRegistry = new ServiceRegistry();
        Map<Class<?>, Object> services = ServiceScanner.scanServices();
        services.forEach((serviceInterface, implementation) -> {
            this.serviceRegistry.registerService(serviceInterface, implementation);
            logger.info("Auto-registered service: {}", serviceInterface.getName());
        });

        if (serverType == ServerType.MASTER) {
            routeManager = new RouteManager();
            startHeartbeatMonitor();
        } else if (serverType == ServerType.WORKER && masterHost != null) {
            masterConnector = new MasterConnector(this, masterHost, masterPort);
        }
    }

    @Override
    public void start() {
        running = true;
        initializeServer();
        startAcceptingConnections();
        logger.info("Cluster server started as {}", serverType);

        if (serverType == ServerType.MASTER) {
            registerSelfAsWorker();
        }

        if (serverType == ServerType.WORKER && masterConnector != null) {
            logger.info("Connecting to master...");
            masterConnector.connect();
        }
    }

    @Override
    public void close() throws IOException {
        logger.info("Closing cluster server...");
        running = false;

        // 关闭后台执行器
        shutdownExecutor(backgroundExecutor);

        // 关闭所有worker连接
        workerConnections.values().forEach(SocketClient::disconnect);
        workerConnections.clear();

        // 关闭主连接器
        if (masterConnector != null) {
            masterConnector.disconnect();
        }

        super.close();
    }

    private void shutdownExecutor(ExecutorService executor) {
        executor.shutdown();
        try {
            if (!executor.awaitTermination(5, TimeUnit.SECONDS)) {
                executor.shutdownNow();
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }

    // ================= 主服务器功能 =================

    private void startHeartbeatMonitor() {
        backgroundExecutor.scheduleAtFixedRate(() -> {
            if (routeManager == null) return;

            long now = System.currentTimeMillis();
            routeManager.getWorkerInfos().values().forEach(worker -> {
                if (!worker.getId().equals(serverId)) {
                    if (now - worker.getLastHeartbeat() > config.getHeartbeatTimeout()) {
                        logger.warn("Worker timeout: {}", worker.getId());
                        routeManager.removeWorker(worker.getId());
                        cleanupWorkerConnection(worker.getId());
                    }
                }
            });
        }, 30, 30, TimeUnit.SECONDS);
    }

    private void registerSelfAsWorker() {
        String host = getPublicIp();
        Set<String> services = this.serviceRegistry.getRegisteredServices();

        // 确保有有效的服务注册
        if (services.isEmpty()) {
            services = Set.of("DEFAULT_SERVICE");
        }

        RouteManager.WorkerInfo selfInfo = new RouteManager.WorkerInfo(
                serverId, host, config.getPort(), System.currentTimeMillis(), services
        );

        routeManager.registerWorker(selfInfo);
        logger.info("Master registered as worker with services:{}", services);
    }

    private String getPublicIp() {
        String host = config.getServerHost();
        if ("localhost".equals(host) || "127.0.0.1".equals(host)) {
            return IpUtil.getLocalIp().orElse(host);
        }
        return host;
    }

    // ================= 从服务器功能 =================

    @Override
    protected void handleConnection(Socket clientSocket) {
        try {
            DataInputStream dataIn = new DataInputStream(clientSocket.getInputStream());
            DataOutputStream dataOut = new DataOutputStream(clientSocket.getOutputStream());

            ClusterHandler handler = new ClusterHandler(this, clientSocket, dataIn, dataOut);
            new Thread(handler).start();
        } catch (IOException e) {
            logger.error("Failed to create connection handler", e);
            closeQuietly(clientSocket);
        }
    }

    private void cleanupWorkerConnection(String workerId) {
        SocketClient client = workerConnections.remove(workerId);
        if (client != null) {
            synchronized (client) {
                client.disconnect();
                logger.info("Worker connection cleaned up: {}", workerId);
            }
        }
    }
}